package dao

import (
	"context"
	"flare-admin/db/sqlc/cluster"
	"flare-admin/db/sqlc/ddos"
	"flare-admin/db/sqlc/micro"
	"flare-admin/db/sqlc/reversesh"
	"flare-admin/db/sqlc/ztrust"
	"fmt"
	"runtime"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

type DBTXQuery struct {
	Micro     *micro.Queries
	ReverseSh *reversesh.Queries
	Cluster   *cluster.Queries
	Ztrust    *ztrust.Queries
	Ddos      *ddos.Queries
	conn      *pgx.Conn
}

func NewDBPool(ctx context.Context, dbHost string) (*pgxpool.Pool, error) {
	cfg, err := pgxpool.ParseConfig(dbHost)
	cfg.MaxConns = int32(runtime.NumCPU())
	if err != nil {
		return nil, err
	}
	pool, err := pgxpool.NewWithConfig(ctx, cfg)
	if err != nil {
		return nil, err
	}

	return pool, nil
}

func NewDBConn(ctx context.Context, dbHost string) (*pgx.Conn, error) {
	config, err := pgx.ParseConfig(dbHost)
	if err != nil {
		return nil, err
	}
	return pgx.ConnectConfig(ctx, config)
}

func NewDBTXQuery(conn *pgx.Conn) *DBTXQuery {
	return &DBTXQuery{
		Micro:     micro.New(conn),
		ReverseSh: reversesh.New(conn),
		Cluster:   cluster.New(conn),
		Ztrust:    ztrust.New(conn),
		Ddos:      ddos.New(conn),
		conn:      conn,
	}
}

func NewDBTXQueryWithConnPool(pool *pgxpool.Pool) *DBTXQuery {
	return &DBTXQuery{
		Micro:     micro.New(pool),
		ReverseSh: reversesh.New(pool),
		Cluster:   cluster.New(pool),
		Ztrust:    ztrust.New(pool),
		Ddos:      ddos.New(pool),
		// conn:      pool.Acquire(),
	}
}

func (txQuery *DBTXQuery) ExecTx(ctx context.Context, isolationLevel pgx.TxIsoLevel, fn func(queries *DBTXQuery) error) error {
	tx, err := txQuery.conn.BeginTx(ctx, pgx.TxOptions{
		IsoLevel: isolationLevel,
	})

	defer func() {
		if p := recover(); p != nil {
			_ = tx.Rollback(ctx) // 避免连接泄露
			panic(p)             // 继续抛出，交给gin层处理这个panic
		}
	}()

	if err != nil {
		return err
	}
	q := NewDBTXQuery(tx.Conn())
	err = fn(q)
	if err != nil {
		if rollbackErr := tx.Rollback(context.Background()); rollbackErr != nil {
			return fmt.Errorf("tx error: %v, rollback err = %v", err, rollbackErr)
		}
		return err
	}
	return tx.Commit(context.Background())
}
